from blackrltask import BlackRLTask
from blackrlenv  import BlackRLEnv
from cantor import Cantor

from pybrain.rl.learners.valuebased import ActionValueTable
from pybrain.rl.agents import LearningAgent
from pybrain.rl.learners import Q
from pybrain.rl.experiments import Experiment
from pybrain.rl.explorers.discrete import EpsilonGreedyExplorer

import numpy
from numpy import genfromtxt

from erroranalysis import ErrorAnalysis
from visualization import vis
import timeit
import sys


if __name__ == "__main__":

    use_counting = False

    if use_counting == False:
        num_states = 2047
    else:
        num_states = 16383

    limit = 1000000

    print "Total possible states: " + str(num_states)
    print "Looping:               " + str(limit)
    print ""

    av_table = ActionValueTable(num_states, 2)
#     av_table.initialize(0.)
    data = genfromtxt("q_matrix.csv", delimiter=",").flatten()
    av_table.initialize(data)

    # define Q-learning agent
    alpha = 1
    gamma = .999
    epsilon = .3
    decay = 0.9999
    
    learner = Q(alpha, gamma)
    explorer = EpsilonGreedyExplorer(epsilon, decay)
    explorer._setModule(av_table)
    learner._setExplorer(explorer)
    agent = LearningAgent(av_table, learner)

    # define the environment
    env = BlackRLEnv()

    # define the task
    task = BlackRLTask(env)

    # finally, define experiment
    experiment = Experiment(task, agent)
    
    # Start time
    start = timeit.default_timer()
    
    # ready to go, start the process
    Err = ErrorAnalysis()
    average = []
    trackbets = []
    interval = limit/100
    percentComplete = 1
    for x in range(0, limit):
        experiment.doInteractions(1)
        agent.learn()
        agent.reset()
        
        table = av_table.params.reshape(num_states,2)
        if x%interval == 0:
            trackbets.append(env.bets.AgentMoney)
            average.append(Err.getAverage(table))
            print(str(percentComplete) + "% Complete...")
            percentComplete = percentComplete + 1

    print ""    
    # End time
    stop = timeit.default_timer()
    
    # Print runtime, and input parameters
    print("runtime:               " + str((stop-start)/60) + " min")
    print("gamma:                 " + str(gamma))
    print("alpha:                 " + str(alpha))
    print("epsilon:               " + str(epsilon))
    print("decay:                 " + str(decay))
    print "Winnings = " + str(env.bets.AgentMoney)
    print ""
    print("Reshaping matrix and writing to file...")
    
    q_matrix = av_table.params.reshape(num_states, 2)
    numpy.savetxt("q_matrix.csv", q_matrix, delimiter=",")
    print("done.")
    
    vis().plot1(range(0,limit,interval),average)
    vis().plot1(range(0,limit,interval),trackbets)
#     import convertqmatrix

